import random

def assign_cluster(x, c):
    min_dist = float('inf')
    cluster_idx = 0
    for i, center in enumerate(c):
        dist = 0
        for a, b in zip(x, center):
            dist += (a - b) ** 2
        if dist < min_dist:
            min_dist = dist
            cluster_idx = i
    return cluster_idx

def Kmeans(data, k, epsilon=1e-6, iteration=100):
    # 随机初始化k个簇中心
    centers = random.sample(data, k)
    n = len(data)
    cluster_labels = [0] * n
    iter_count = 0

    while iter_count < iteration:
        # 为每个样本分配簇
        new_cluster_labels = [assign_cluster(x, centers) for x in data]
        # 更新簇中心
        new_centers = []
        for i in range(k):
            cluster_samples = [data[j] for j in range(n) if new_cluster_labels[j] == i]
            if not cluster_samples:
                # 如果簇为空，重新随机选一个样本作为中心
                new_center = random.choice(data)
            else:
                # 计算簇内样本的均值作为新中心
                dim = len(data[0])
                new_center = [0.0] * dim
                for sample in cluster_samples:
                    for d in range(dim):
                        new_center[d] += sample[d]
                for d in range(dim):
                    new_center[d] /= len(cluster_samples)
            new_centers.append(new_center)
        # 检查簇中心变化是否小于阈值
        center_change = 0
        for c1, c2 in zip(centers, new_centers):
            for a, b in zip(c1, c2):
                center_change += (a - b) ** 2
        if center_change < epsilon:
            break
        centers = new_centers
        cluster_labels = new_cluster_labels
        iter_count += 1

    return centers, cluster_labels

if __name__ == "__main__":
    data = [
        [1, 2], [1, 4], [1, 0],
        [4, 2], [4, 4], [4, 0],
    ]
    k = 2
    centers, labels = Kmeans(data, k)
    print("簇中心:", centers)
    print("样本簇标签:", labels)